from test_unet_keras import get_file_list, img_size
from UNetPlusPlus.keras.segmentation_models import Unet
from utils.helper_functions import mean_iou_numpy
import numpy as np
import os
import skimage.io as skio
import tensorflow as tf
from skimage.util import img_as_ubyte
from skimage.transform import rotate
from tqdm import tqdm
from tensorflow.keras import backend as K
from tensorflow.keras.losses import CategoricalCrossentropy
from tensorflow.keras.optimizers import Adam
from skimage.transform import resize
import re
import csv
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, classification_report

os.environ["CUDA_VISIBLE_DEVICES"] = "2"
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)


image_path = r"/home/zyp/project/EI_SEG/dataset/cherry_1024/img/"
mask_path = r"/home/zyp/project/EI_SEG/dataset/cherry_1024/mask/"


def load_data(file_list, isMask=False):
    img_list = []
    for f in tqdm(file_list):
        img = skio.imread(f)
        img = resize(img, img_size)
        if isMask:
            major_vein = img[:, :, 0] > 0
            sec_vein = img[:, :, 1] > 0
            path_strs = os.path.split(f)
            leaf_name = re.findall("\d{1,}\S{1,}_mask", path_strs[-1])[0][0:-5]
            leaf_img = skio.imread(os.path.join(image_path, leaf_name + '.png'))
            leaf_img = resize(leaf_img, img_size)
            leaf_mask = leaf_img[:, :, 0] > 0
            mask = ~(major_vein | sec_vein)
            leaf_part = leaf_mask & mask
            background = mask & ~leaf_part
            img = np.dstack([major_vein.astype(np.uint8), sec_vein.astype(np.uint8), leaf_part.astype(np.uint8), background.astype(np.uint8)])
        img_list.append(img)
    return np.asarray(img_list)


def compute_iou(y_true, y_pred):
    y_true = y_true[:, :, 0:3]
    y_pred = y_pred[:, :, 0:3]
    y_true_num = np.argmax(y_true, axis=-1)
    y_pred_num = np.argmax(y_pred, axis=-1)
    y_true_label = y_true_num.flatten(order='C')
    y_pred_label = y_pred_num.flatten(order='C')
    cm = confusion_matrix(y_true_label, y_pred_label)
    report = classification_report(y_true_label, y_pred_label, output_dict=True)
    iou_major = cm[0][0] / (cm[0][0] + cm[1][0]+cm[0][1])
    iou_sec = cm[1][1] / (cm[1][1] + cm[0][1] + cm[1][0])
    print(report)
    return {'major_iou': iou_major,
            'sec_iou': iou_sec,
            'report': report}


if __name__ == "__main__":
    cherry_img_file_list = get_file_list(image_path)
    for index, f in enumerate(cherry_img_file_list):
        cherry_img_file_list[index] = os.path.join(image_path, f)

    imgs = load_data(cherry_img_file_list[0:1])

    cherry_mask_file_list = get_file_list(mask_path)
    for index, f in enumerate(cherry_mask_file_list):
        cherry_mask_file_list[index] = os.path.join(mask_path, f)

    masks = load_data(cherry_mask_file_list[0:1], isMask=True)

    net = Unet(backbone_name='resnet50',
               input_shape=img_size,
               classes=4,
               encoder_weights='imagenet',
               activation='softmax')

    net.summary()
    #loss = weightedLoss(tf.keras.losses.categorical_crossentropy, class_weights)
    loss = CategoricalCrossentropy(from_logits=True)
    adam = Adam(learning_rate=0.001, name='adam')
    net.compile(optimizer=adam, loss=loss, metrics=["categorical_accuracy"])
    net.load_weights(r"./unet-resnet50.h5")
    iou_results = []
    file_names = []
    predict_result = {}
    for index, img in enumerate(tqdm(imgs)):
        img = np.reshape(img, [1, 512, 512, 3])
        pred = net.predict(img)
        mask = masks[index]
        pred = np.reshape(pred, [512, 512, 4])
        #mean_iou = mean_iou_numpy(mask, pred)
        metric_result = compute_iou(mask, pred)
        result_img = np.column_stack([mask[:, :, 0:3], pred[:, :, 0:3]])
        file_names.append(cherry_img_file_list[index])

        f_name = os.path.split(cherry_img_file_list[index])[-1]
        skio.imsave("./cherry_pred/{}_{}_pred.png".format(index, f_name[0:-4]), result_img)
        predict_result[f_name[0:-4]] = metric_result

